import pandas as pd
import sqlite3
import requests
from typing import Union, Dict, Any
from pathlib import Path

class DataLoader:
    """支持多种数据源的数据加载器"""
    
    def __init__(self):
        self.supported_sources = {
            'excel': self._load_excel,
            'csv': self._load_csv,
            'json': self._load_json,
            'sqlite': self._load_sqlite,
            'api': self._load_from_api
        }
    
    def load(self, source: str, **kwargs) -> pd.DataFrame:
        """加载数据源"""
        source_type = self._detect_source_type(source)
        if source_type not in self.supported_sources:
            raise ValueError(f"不支持的源类型: {source_type}")
        
        return self.supported_sources[source_type](source, **kwargs)
    
    def _detect_source_type(self, source: str) -> str:
        """自动检测源类型"""
        if source.startswith(('http://', 'https://')):
            return 'api'
        elif source.endswith('.xlsx'):
            return 'excel'
        elif source.endswith('.csv'):
            return 'csv'
        elif source.endswith('.json'):
            return 'json'
        elif source.endswith('.db') or ':memory:' in source:
            return 'sqlite'
        else:
            return 'excel'  # 默认尝试Excel
    
    def _load_excel(self, file_path: str, **kwargs) -> pd.DataFrame:
        """加载Excel文件"""
        sheet_name = kwargs.get('sheet_name', 0)
        return pd.read_excel(file_path, sheet_name=sheet_name)
    
    def _load_csv(self, file_path: str, **kwargs) -> pd.DataFrame:
        """加载CSV文件"""
        return pd.read_csv(file_path, **kwargs)
    
    def _load_json(self, file_path: str, **kwargs) -> pd.DataFrame:
        """加载JSON文件"""
        return pd.read_json(file_path, **kwargs)
    
    def _load_sqlite(self, db_path: str, **kwargs) -> pd.DataFrame:
        """从SQLite数据库加载数据"""
        table = kwargs.get('table')
        query = kwargs.get('query')
        
        if not table and not query:
            raise ValueError("必须指定表名或查询语句")
        
        conn = sqlite3.connect(db_path)
        
        if query:
            return pd.read_sql(query, conn)
        else:
            return pd.read_sql(f"SELECT * FROM {table}", conn)
    
    def _load_from_api(self, url: str, **kwargs) -> pd.DataFrame:
        """从API加载数据"""
        params = kwargs.get('params', {})
        response = requests.get(url, params=params)
        response.raise_for_status()
        
        data = response.json()
        if isinstance(data, list):
            return pd.DataFrame(data)
        elif isinstance(data, dict) and 'data' in data:
            return pd.DataFrame(data['data'])
        else:
            raise ValueError("API返回的数据格式不支持")